import numpy as np
from .quanternion import *


def export_ADE_along_time(pred, gt, pre):
    diff = pred - gt
    norm = torch.linalg.norm(diff, axis=3).mean(axis=2)
    ADE_along_time = norm.mean(axis=0) / 100
    ADE_along_time = np.asarray(ADE_along_time.cpu())

    # np.savetxt(pre+'@ADE_along_time.txt', ADE_along_time)

    return ADE_along_time.reshape(ADE_along_time.shape[0], 1)


def cal_ADE_FDE_metric(pred, gt):
    """
    Average displacement error & Final displacement error
    :param pred: pred global position, BxTxJx3
    :param gt: gt global position
    :return: ADE, FDE
    """
    diff = pred - gt
    # norm = torch.linalg.norm(diff, ord=2, axis=3).mean(axis=2)
    norm = torch.linalg.norm(diff, axis=3).mean(axis=2)
    ADE = norm.mean(axis=1).mean() / 100
    FDE = norm[:, -1].mean() / 100
    # FDE = norm[:, -2].mean() / 100

    return np.asarray(ADE.cpu()), np.asarray(FDE.cpu())


def cal_F5DE_metric(pred, gt):
    diff = pred - gt
    norm = torch.linalg.norm(diff, ord=2, axis=3).mean(axis=2)
    F5DE = norm[:, -5:].mean() / 100

    return np.asarray(F5DE.cpu())


def cal_APD(all_positions):
    length = len(all_positions)
    distances = 0
    for i in range(length):
        for j in range(length):
            if j != i:
                diff = all_positions[i] - all_positions[j]
                dis = torch.linalg.norm(diff, axis=3).sum(axis=2).sum(1).mean() / 100
                distances += dis
    return distances / (length * (length - 1))
